package com.wit.witmg.util;

import java.io.*;
import java.net.*;

public class ModbusTCPClient {
    private Socket socket;
    private DataInputStream input;
    private DataOutputStream output;
    private String host;
    private int port;
    private int transactionId = 0;

    public ModbusTCPClient(String host, int port) {
        this.host = host;
        this.port = port;
    }

    public boolean connect() {
        try {
            socket = new Socket(host, port);
            input = new DataInputStream(socket.getInputStream());
            output = new DataOutputStream(socket.getOutputStream());
            return true;
        } catch (IOException e) {
            System.err.println("连接失败: " + e.getMessage());
            return false;
        }
    }

    public void disconnect() {
        try {
            if (input != null) input.close();
            if (output != null) output.close();
            if (socket != null) socket.close();
        } catch (IOException e) {
            System.err.println("断开连接错误: " + e.getMessage());
        }
    }

    public byte[] createReadRequest(int unitId, int functionCode, int startAddress, int quantity) {
        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        DataOutputStream dos = new DataOutputStream(baos);

        try {
            // MBAP头
            dos.writeShort(transactionId++);  // 事务标识符
            dos.writeShort(0x0000);           // 协议标识符
            dos.writeShort(0x0006);           // 长度(后续字节数)
            dos.writeByte(unitId);            // 单元标识符

            // Modbus PDU
            dos.writeByte(functionCode);      // 功能码
            dos.writeShort(startAddress - 1); // 起始地址(协议地址-1)
            dos.writeShort(quantity);         // 读取数量

        } catch (IOException e) {
            e.printStackTrace();
        }

        return baos.toByteArray();
    }

    // 写入单个寄存器请求 (功能码0x06)
    public byte[] createWriteSingleRegisterRequest(int unitId, int address, int value) {
        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        DataOutputStream dos = new DataOutputStream(baos);

        try {
            // MBAP头
            dos.writeShort(transactionId++);
            dos.writeShort(0x0000);
            dos.writeShort(0x0006);           // 长度
            dos.writeByte(unitId);

            // Modbus PDU
            dos.writeByte(0x06);              // 功能码
            dos.writeShort(address - 1);      // 寄存器地址
            dos.writeShort(value);            // 写入值

        } catch (IOException e) {
            e.printStackTrace();
        }

        return baos.toByteArray();
    }

    // 写入多个寄存器请求 (功能码0x10)
    public byte[] createWriteMultipleRegistersRequest(int unitId, int startAddress, int[] values) {
        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        DataOutputStream dos = new DataOutputStream(baos);

        try {
            int quantity = values.length;
            int byteCount = quantity * 2;

            // MBAP头
            dos.writeShort(transactionId++);
            dos.writeShort(0x0000);
            dos.writeShort(7 + byteCount);    // 长度: 7 + 数据字节数
            dos.writeByte(unitId);

            // Modbus PDU
            dos.writeByte(0x10);              // 功能码
            dos.writeShort(startAddress - 1); // 起始地址
            dos.writeShort(quantity);         // 寄存器数量
            dos.writeByte(byteCount);         // 字节数

            // 写入数据
            for (int value : values) {
                dos.writeShort(value);
            }

        } catch (IOException e) {
            e.printStackTrace();
        }

        return baos.toByteArray();
    }

    public byte[] sendRequest(byte[] request) throws IOException {
        output.write(request);
        output.flush();

        // 读取MBAP头
        byte[] header = new byte[7];
        input.readFully(header);

        // 获取数据长度
        int length = ((header[4] & 0xFF) << 8) | (header[5] & 0xFF);

        // 读取数据部分
        byte[] data = new byte[length - 1]; // 减去单元标识符
        input.readFully(data);

        return data;
    }
}
